import traceback

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.cluster import HDBSCAN, DBSCAN
from sklearn.decomposition import PCA

from src.bean.x_typing import PRE_HANDLE_TYPE
from src.util.common_util import printx, while_input, is_int_between


def plot_cluster(df, labels, out_path="cluster_plot.png", title="result"):
    pca = PCA(n_components=2)
    points_2d = pca.fit_transform(df.values)

    labels = np.array(labels)
    unique_labels = set(labels)
    colors = plt.cm.get_cmap("tab10", len(unique_labels))

    plt.figure(figsize=(10, 6))
    for i, label in enumerate(unique_labels):
        mask = labels == label
        if label == -1:
            plt.scatter(points_2d[mask, 0], points_2d[mask, 1], c='red', label='noise', marker=".")
        else:
            plt.scatter(points_2d[mask, 0], points_2d[mask, 1], c=[colors(i)], label=f'class {label}')

    plt.title(title)
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig(out_path)
    plt.close()


def scan(df: pd.DataFrame, pre_handle_type: PRE_HANDLE_TYPE) -> (pd.DataFrame, str):
    if len(df) == 0:
        return df.copy(), "数据为空，没有聚类"
    df = df.reset_index(drop=True).copy()
    try:
        if pre_handle_type == PRE_HANDLE_TYPE.train or pre_handle_type == PRE_HANDLE_TYPE.increase:
            scan_type = int(
                while_input(f"[1:hdbscan(效果较好，慢，建议用这个)|2:dbscan|3:不降噪]:", is_int_between,
                            (1, 3)))
            if scan_type == 1:
                cluster_str = "HDBSCAN"
                cluster = HDBSCAN(n_jobs=36)
            elif scan_type == 2:
                cluster_str = "DBSCAN"
                cluster = DBSCAN(min_samples=len(df.columns) * 2, eps=21.18)  # 去掉 n_jobs
            else:
                return df.copy(), "没有选择聚类"
            cluster.fit(df.values)
            labels = cluster.labels_
            if len(labels) != len(df):
                printx(f"[预处理]警告：labels数量({len(labels)})与df行数({len(df)})不符，自动裁切对齐")
                min_len = min(len(labels), len(df))
                labels = labels[:min_len]
                df = df.iloc[:min_len].copy()
            noise_mask = labels == -1
            noise_count = np.sum(noise_mask)
            printx(f"[预处理]噪点数量:{noise_count}/{len(labels)} ({(noise_count / len(labels) * 100):.1f}%)")
            result = df[~noise_mask].copy()
            printx(f"[预处理]最终数量:{len(result)} (保留率：{(len(result) / len(df) * 100):.1f}%)")
            return result.reset_index(drop=True), cluster_str
        else:
            return df, "非训练模式，没有聚类"
    except Exception as e:
        traceback.print_exc()
        printx(f"[预处理]聚类失败：{str(e)}")
        return df.copy(), f"聚类失败：{str(e)}"
